# -*- coding: utf-8 -*-
'''
Created on 12.07.2019

@author: yu03
'''
from FFT_Interpolation import *
import cv2
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
from scipy.interpolate import interp1d


def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
#     return np.concatenate((np.zeros(N-1), (cumsum[N:] - cumsum[:-N]) / float(N)))
    return np.concatenate((np.zeros(int((N-1)/2)), (cumsum[N:] - cumsum[:-N]) / float(N), np.zeros(int((N)/2))))

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
#     return a*np.exp(-(x-b)**2/2/c**2)

p= -0.021335
# p= -0.02129
def fit_func_phase(x, p, q):
    return p*x+q

# if 0: ### Get Real Noise_PSD
# #     img_set = []
#     pattern_path = r'C:\Users\yu03\Desktop\Noise_50fps_180offset_20usExpo.avi'
#     cap = cv2.VideoCapture(pattern_path)
#     frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
#     print(frame_num, "Frames")
#     # while(cap.isOpened()):
#     for k in range(frame_num):
#         ret, frame = cap.read() 
#         img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# #         img_set.append(img)
#     line_noise = img[512]
#     plt.figure('real line noise')
#     plt.plot(line_noise)
#     plt.show()

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

# pix_size = 5.3e-6
pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size

x = np.linspace(0, (pix_num-1)*pix_size, pix_num)
tau0 = pix_size
fs = 1/tau0
# x = np.linspace(0, 1, 1001)
# fs = 1000
# tau0 = 1 / fs
N = len(x)

frequency_estim_set_sim = []
frequency_estim_set_1 = []
frequency_estim_set_2 = []
frequency_estim_set_3 = []
frequency_estim_set_4 = []

phase_estim_set_sim = []
phase_estim_set_1 = []
phase_estim_set_2 = []
phase_estim_set_3 = []
phase_estim_set_4 = []
freq_for_phase_set = []
f_fit_set = []
f_fit_np_set = []
phi_fit_set = []

m_k_num_set = []
X_m_k_set = []

f = np.linspace(5866, 5866, 500)
phi = np.linspace(0, 0*np.pi, 500)
phi = []
for i in range(10):
    for j in range(50):
        phi.append([i*1/50/25]) ### 50nm/rad
# print(phi)
print(len(phi))
# f = np.linspace(4000, 8000, 2000)
# phi = np.linspace(0, 0, 2000)

for i in range(len(f)):
    sig = 1024*np.cos(2*np.pi*f[i]*x + phi[i])/2 + 1024/2
#     print(i, sig[0], sig[500])
#     sig += 0.5*np.cos(2*np.pi*300*x + 0)
    window = signal.gaussian(len(sig), std=len(sig)/10)
#     window = signal.hann(len(sig))
#     if 0:
#         shift_num = 200
#         gaussian_base = signal.gaussian(len(sig)-shift_num, std=(len(sig)-shift_num)/10)
#         gaussian_1 = np.concatenate((gaussian_base, np.zeros(shift_num)))
#         gaussian_2 = np.concatenate((np.zeros(shift_num), gaussian_base))
#         window = (gaussian_1 + gaussian_2)/2
    sig *= window
#     line_noise = np.random.normal(0, 0.6, len(sig)) 
    line_noise = np.random.normal(2.43, 0.8, len(sig)) 
    line_noise = line_noise.round().astype(int)
    sig_sim = sig
    sig = sig + line_noise
    
#     interplation_func = interp1d(x, sig_sim, kind='cubic')
#     x_interplation = np.linspace(0, (pix_num-1)*pix_size, pix_num*10)
#     sig_sim = interplation_func(x_interplation)
#        
#     interplation_func = interp1d(x, sig, kind='cubic')
#     x_interplation = np.linspace(0, (pix_num-1)*pix_size, pix_num*10)
#     sig = interplation_func(x_interplation)
#        
#     tau0 = pix_size/10
    
#     sig = np.diff(sig)
#     sig = np.concatenate(([0], sig))
#     sig_sim = np.diff(sig_sim)
#     sig_sim = np.concatenate(([0], sig_sim))
    
#     sig = sig - running_mean(sig, 50)
#     sig_sim = sig_sim - running_mean(sig_sim, 50)
#     print(len(sig))

    DC_num = 1000
#     freq_estim_1, phase_estim_1 = FFT_interpolation_1(sig, tau0, 1e5, 3000)
    freq_estim_sim, phase_estim_sim, freqline_sim, sig_magnitude_sim, sig_phase_sim,  m_k_num_sim, X_m_k_sim, freq_for_phase_sim  = FFT_interpolation_2(sig_sim, tau0, 1e5, DC_num)
    freq_estim_2, phase_estim_2, freqline, sig_magnitude, sig_phase,  m_k_num, X_m_k, freq_for_phase  = FFT_interpolation_2(sig, tau0, 1e5, DC_num)
#     freq_estim_3, phase_estim_3 = FFT_interpolation_3(sig, tau0, 1e5, 3000)
#     freq_estim_4, phase_estim_4 = FFT_interpolation_4(sig, tau0, 1e5, 3000)
    
#     frequency_estim_set_1.append(freq_estim_1-f[i])
    frequency_estim_set_2.append(freq_estim_2-f[i])
    frequency_estim_set_sim.append(freq_estim_sim)
#     frequency_estim_set_3.append(freq_estim_3-f[i])
#     frequency_estim_set_4.append(freq_estim_4-f[i])
    
#     phase_estim_set_1.append(phase_estim_1*2-np.pi)
#     phase_estim_set_2.append(phase_estim_2*2-np.pi)
    phase_estim_set_2.append(phase_estim_2)
    phase_estim_set_sim.append(phase_estim_sim)
#     phase_estim_set_3.append(phase_estim_3*2-np.pi)
#     phase_estim_set_4.append(phase_estim_4*2-np.pi)
    print(i, m_k_num, X_m_k, freqline[m_k_num])
    m_k_num_set.append(m_k_num)
    X_m_k_set.append(X_m_k)
    freq_for_phase_set.append(freq_for_phase)
    FFT_start = np.where(sig_magnitude[DC_num:] > X_m_k*0.4)[0][0]+DC_num
    FFT_end = np.where(sig_magnitude[DC_num:] > X_m_k*0.4)[0][-1]+DC_num
    fit_x = freqline[FFT_start:FFT_end+1]
    fit_freq = sig_magnitude[FFT_start:FFT_end+1]
    fit_phase = sig_phase[FFT_start:FFT_end+1]
    fit_x_sim = freqline_sim[FFT_start:FFT_end+1]
    fit_freq_sim = sig_magnitude_sim[FFT_start:FFT_end+1]
    fit_phase_sim = sig_phase_sim[FFT_start:FFT_end+1]
    
    params = curve_fit(fit_func, np.array(fit_x), np.array(fit_freq))
    [a, b, c] = params[0]
    f_fit = b
    f_fit_set.append(f_fit-f[i])
    
#     [aa, bb, cc] = np.polyfit(fit_x, fit_freq, 2)
#     f_fit_np = -bb/2/aa
#     f_fit_np_set.append(f_fit_np)
    
    params_phase = curve_fit(fit_func_phase, fit_x, np.unwrap(fit_phase))
    [p, q] = params_phase[0]
#     p= -0.021314
    phi_fit = p*freqline[m_k_num]+q
    phi_fit_set.append(phi_fit+2*np.pi)
#     print(p,q)
#     print(f_fit, phi_fit, sig_phase[m_k_num])
    
    if i == int(len(f)-1):
#     if i < 400:
        plt.figure('Simulated Line Signal')
        plt.subplot(5,1,1)
#         plt.plot(gaussian_1, 'b')
#         plt.plot(gaussian_2, 'k')
        plt.plot(window, 'r')        
        plt.title("Signal")
        plt.ylabel("Amplitude")
        plt.xlabel("Samples")
        plt.grid(which='major', axis='both')
         
        plt.subplot(5,1,2)
        plt.plot(sig, 'b')
        plt.plot(sig_sim, 'r')
        plt.title("Signal")
        plt.ylabel("Amplitude")
        plt.xlabel("Samples")
        plt.grid(which='major', axis='both')
         
        plt.subplot(5,1,3)
        plt.plot(freqline, sig_magnitude, 'b')
        plt.plot(freqline_sim, sig_magnitude_sim, 'r')
#         plt.stem(freqline, sig_magnitude, linefmt='b', markerfmt='bo', use_line_collection='True')
#         plt.stem(freqline_sim, sig_magnitude_sim, linefmt='r', markerfmt='ro', use_line_collection='True')
        plt.grid(which='major', axis='both')
         
        plt.subplot(5,1,4)
        plt.plot(m_k_num_set)
        plt.grid(which='major', axis='both')
         
        plt.subplot(5,1,5)
        plt.plot(X_m_k_set)
        plt.grid(which='major', axis='both')
        
        
        plt.figure('FFT fit')
        plt.subplot(3,2,1)
        plt.plot(freqline[2150:4150], sig_magnitude[2150:4150], color='b', marker='')
#         plt.plot(fit_x, a*(fit_x-b)**2+c, 'r')
        plt.grid(which='major', axis='both')
#         plt.subplot(3,2,3)
#         plt.plot(fit_x, fit_freq, color ='b', marker='o')
#         plt.plot(fit_x, fit_freq_sim, color ='r', marker=' ')
#         plt.plot(fit_x, a*(fit_x-b)**2+c, 'k')
#         plt.grid(which='major', axis='both')
#         plt.subplot(3,2,5)
#         plt.plot(fit_x, fit_freq-fit_freq_sim, marker=' ')
# #         plt.plot(fit_x, a*(fit_x-b)**2+c, 'r')
#         plt.grid(which='major', axis='both')

        plt.subplot(3,2,3)
#         plt.stem(freqline, sig_phase, use_line_collection='True')
#         plt.plot(freqline_sim, sig_phase_sim, 'r')
        plt.plot(freqline[2850:3450], sig_phase[2850:3450], 'b')
        plt.grid(which='major', axis='both')
#         plt.subplot(3,2,4)
#         plt.plot(fit_x, np.unwrap(fit_phase), marker='o')
#         plt.plot(fit_x, np.unwrap(fit_phase_sim), color='r', marker=' ')
#         plt.plot(fit_x, p*(fit_x)+q, 'k')
#         plt.grid(which='major', axis='both')
#         plt.subplot(3,2,6)
#         plt.plot(fit_x, np.unwrap(fit_phase)-np.unwrap(fit_phase_sim), 'b', marker=' ')
#         phase_fit_error = p*fit_x+q-np.unwrap(fit_phase_sim)
#         plt.plot(fit_x, phase_fit_error, 'k', marker=' ')
#         plt.grid(which='major', axis='both')
# #         plt.ylim(-0.0025,0.0025)
# #         plt.show()
plt.figure('Estimation Compare')
# plt.subplot(4,2,1)
# plt.plot(frequency_estim_set_1)
# plt.title("Eq.1")
# plt.ylabel("Amplitude")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')
# plt.subplot(4,2,2)
# plt.plot(np.unwrap(phase_estim_set_1)/2+np.pi/2-phi)
# plt.title("Eq.1")
# plt.ylabel("rad")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')

ax3 = plt.subplot(4,2,3)
# plt.plot(freq_for_phase_set-f, 'r')
plt.plot(frequency_estim_set_2, 'b')
# plt.plot(f_fit_np_set-f, 'r')
plt.plot(np.array(f_fit_set), 'k', marker='o', markersize=2)

# plt.plot(frequency_estim_set_sim-f, 'r')

plt.title("Eq.2")
plt.ylabel("Frequency Estimation [/m]")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
V_x = 0
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax3_angle = ax3.twinx()
plt.ylim(angle_range)
plt.ylabel("Tilting [urad]")

ax4 = plt.subplot(4,2,4)
# plt.plot(np.unwrap(phase_estim_set_2)/2+np.pi/2-phi)

plt.plot(np.unwrap(phase_estim_set_2), 'b')
# plt.plot(np.unwrap(phi_fit_set)-np.unwrap(phi), 'k')
plt.plot(np.unwrap(phi_fit_set), 'k', marker='o', markersize=2)


# plt.plot(np.array(freq_for_phase_set-f)/100, 'r')

plt.title("Eq.2")
plt.ylabel("phase [rad]")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax4_length = ax4.twinx()
plt.ylim(length_range)
plt.ylabel('Length (nm)')

# plt.subplot(4,2,5)
# plt.plot(frequency_estim_set_3)
# plt.title("Eq.3")
# plt.ylabel("Amplitude")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')
# plt.subplot(4,2,6)
# # plt.plot(np.unwrap(phase_estim_set_3)-phi_set)
# plt.plot(np.unwrap(phase_estim_set_3)/2+np.pi/2-phi)
# plt.title("Eq.4")
# plt.ylabel("rad")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')
# 
# plt.subplot(4,2,7)
# plt.plot(frequency_estim_set_4)
# plt.title("Eq.4")
# plt.ylabel("Amplitude")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')
# plt.subplot(4,2,8)
# # plt.plot(np.unwrap(phase_estim_set_4)-phi_set)
# plt.plot(np.unwrap(phase_estim_set_4)/2+np.pi/2-phi)
# plt.title("Eq.4")
# plt.ylabel("rad")
# plt.xlabel("Samples")
# plt.grid(which='major', axis='both')

plt.tight_layout()

plt.show()
